#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar 29 21:02:40 2024

@author: lixiao
"""
 
import numpy as np
from scipy.spatial.transform import Rotation as sR
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection


def generateCube(n, a):
    delta = a / (n - 1)
    pts = []
    for i in np.arange(-a / 2, a / 2 + delta, delta):
        for j in np.arange(-a / 2, a / 2 + delta, delta):
            for k in np.arange(-a / 2, a / 2 + delta, delta):
                if abs(i) == a / 2 or abs(j) == a / 2 or abs(k) == a / 2:
                    pts.append([i, j, k])
    return np.array(pts)


def generateCameraTrajectory(numViews, radius):
    """
    Generates a circular camera trajectory about the origin, 
    assuming the radius of the trajectory is `radius`,
    and that all views (a total of `numViews`) are equally spaced along the circle. 
    The first point on the trajectory is assumed to lie on the X-axis.

    :param numViews: Number of views (points on the trajectory)
    :param radius: Radius of the circular trajectory
    :return Rs: List of rotation matrices for each view
    :return ts: List of translation vectors for each view
    """
    # Angle increment at each time step
    angle_increment = 2 * np.pi / numViews

    # Variables to hold the rotation matrices and translation vectors
    Rs = []
    ts = []

    # Generate each trajectory point
    for i in range(1, numViews + 1):
        # Current camera yaw angle (about Z-axis in world frame)
        cur_yaw_angle = angle_increment * (i - 1)

        # Rotation matrix (first, rotate about the world Y-axis by -pi/2 and
        # then about the obtained frame's Z-axis by pi/2, and then about the
        # thus obtained frame's Y-axis by -1*curYawAngle
        R = sR.from_euler('x', -np.pi/2) * sR.from_euler('y', -np.pi/2 - cur_yaw_angle)
        
        # Translation vector
        ts.append(np.array([radius * np.cos(cur_yaw_angle),\
                            radius * np.sin(cur_yaw_angle), 0]))

        # Store the rotation matrix in the list
        Rs.append(R.as_matrix())

    return np.array(Rs), np.array(ts)

def draw_camera(ax, extrinsincs, scale = 1, factor = 2):
  v = np.array([[-0.5, -0.5, 1],
                [0.5, -0.5, 1],
                [0.5, 0.5, 1],
                [-0.5, 0.5, 1],
                [0, 0, 0]])
  v = v * factor
  v = np.c_[v, np.ones(v.shape[0])]
  for E in extrinsincs:
      E = np.vstack((E, np.array([0., 0., 0., 1.])))
      # E[0:3,0:3] = np.linalg.inv(E[0:3,0:3])
      # E[0:3,3] *= -1
      E_inv = E
      E_inv = E_inv[:3]
      
      v_new = np.dot(v, E_inv.T)
      verts = [[v_new[0], v_new[1], v_new[4]], [v_new[0], v_new[3], v_new[4]],
               [v_new[2], v_new[1], v_new[4]], [v_new[2], v_new[3], v_new[4]],
               [v_new[0], v_new[1], v_new[2], v_new[3]]]
      ax.add_collection3d(Poly3DCollection(verts, facecolors='cyan',\
                          linewidths=0.5, edgecolors='r', alpha=.25))

def plotCameraFrustum(ax, R, camLoc):
    # Size unit of the frustum
    frustumSize = 1
    
    # Vertices of the frustum
    frustumVertices = np.array([
        [0, 0, 0],
        [-frustumSize, -0.5*frustumSize, frustumSize],
        [frustumSize, -0.5*frustumSize, frustumSize],
        [frustumSize, 0.5*frustumSize, frustumSize],
        [-1*frustumSize, 0.5*frustumSize, frustumSize]
    ])
    
    # Convert the frustum to world coordinates
    frustumVertices = np.dot(R, frustumVertices.T) + np.tile(camLoc, (5, 1)).T
    
    # Define the edges of the frustum
    frustumEdges = np.array([
        [0, 1], [0, 2], [0, 3], [0, 4],
        [1, 2], [2, 3], [3, 4], [4, 1]
    ])
    
    for edge in frustumEdges:
        ax.plot3D(frustumVertices[0, edge], frustumVertices[1, edge],\
                  frustumVertices[2, edge], 'b-')

def synthesizeImages(X, K, Rs, ts):
    """
    Synthesizes images from 3D points using camera parameters.
    Parameters:
        X : numpy.ndarray
            3 x N array of N 3D points.
        K : numpy.ndarray
            3 x 3 camera intrinsics matrix.
        Rs : numpy.ndarray
            M x 3 x 3 array of rotation matrices, where M is the number of views.
        ts : numpy.ndarray
            M x 3 array of translation vectors.
    Returns:
        images : list
            A list of M images, where each image is a 3 x N array containing
            the x and y image coordinates of the projections of each 3D point.
    Notes:
        The function assumes the convention X_cam = R^T * (X_world - t) 
        while converting from world to camera coordinates.
    """
    # Number of views
    numViews = ts.shape[0]
    # Number of points
    numPoints = X.shape[0]
    # List to store images
    images = []
    # Synthesize the images
    for i in range(numViews):
        images_homogeneous = np.dot(K, np.dot(np.transpose(Rs[i,:,:]), \
                                    (X - np.tile(ts[i,:], (numPoints,1))).T))
        images_homogeneous = images_homogeneous \
                            / np.tile(images_homogeneous[2,:], (3,1))
        images.append(images_homogeneous)
    return np.array(images)


''' main '''
if __name__ == '__main__':
  np.set_printoptions(suppress=True)
  
  f = 330
  alpha_ccd = 1
  cx = 240
  cy = 320
  K = np.array([[alpha_ccd*f, 0, cx],
               [0, alpha_ccd*f, cy],
               [0, 0, 1]])
  
  world_points_gt = generateCube(4, 16)
  num_views = 8
  num_points = world_points_gt.shape[0]
  Rs, ts = generateCameraTrajectory(num_views, 25)
  cam_locs = ts.T
  Ps = np.zeros((num_views, 3, 4))
  
  fig = plt.figure()
  ax = fig.add_subplot(111, projection='3d')
  
  ax.scatter(world_points_gt[:, 0], world_points_gt[:, 1], world_points_gt[:, 2], 'filled')
  ax.set_xlabel('X')
  ax.set_ylabel('Y')
  ax.set_zlabel('Z')
  ax.set_xlim(-26,26)
  ax.set_ylim(-26,26)
  ax.set_zlim(-26,26)
  ax.set_title('The synthetic scene and cameras')
  
  ex = []
  for i in range(num_views):
      R_cam_W = Rs[i]
      ex.append(np.c_[R_cam_W, cam_locs[:, i]])
  draw_camera(ax, ex)
  plt.show()
  
  images = synthesizeImages(world_points_gt, K, Rs, ts)
  
  fig, axes = plt.subplots(2, 4, figsize=(12, 8))  
  fig.subplots_adjust(hspace=0.4, wspace=0.3)  
  for i in range(num_views):
      row = i // 4  
      col = i % 4  
      ax = axes[row, col]  
      ax.scatter(images[i][0, :], images[i][1, :], marker='o')
      ax.set_xlabel('x')
      ax.set_ylabel('y')
      ax.set_aspect('equal')
      ax.set_title('Image {}'.format(i+1))
  plt.tight_layout()
  plt.show()













